!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep

import grpc

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

SWITCH_TO_HOST_PORT = 1

def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr):

    if tunnel_id == 100:
        SWITCH_TO_SWITCH_PORT = 2
    elif tunnel_id == 101:
        SWITCH_TO_SWITCH_PORT = 2
    elif tunnel_id == 200:
        SWITCH_TO_SWITCH_PORT = 3
    elif tunnel_id == 201:
        SWITCH_TO_SWITCH_PORT = 2
    elif tunnel_id == 300:
        SWITCH_TO_SWITCH_PORT = 3
    elif tunnel_id == 301:
        SWITCH_TO_SWITCH_PORT = 3

    # 1) Tunnel Ingress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.myTunnel_ingress",
        action_params={
            "dst_id": tunnel_id,
        })
    ingress_sw.WriteTableEntry(table_entry)

    # 2) Tunnel Transit Rule on Ingress Switch
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_forward",
        action_params={
            "port": SWITCH_TO_SWITCH_PORT
        })
    ingress_sw.WriteTableEntry(table_entry)

    # 3) Tunnel Egress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_egress",
        action_params={
            "dstAddr": dst_eth_addr,
            "port": SWITCH_TO_HOST_PORT
        })
    egress_sw.WriteTableEntry(table_entry)

def writeArp_matchRules(p4info_helper, sw, match_fields_1, match_fields_2, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.arp_match",
        match_fields={
            "hdr.arp.oper": match_fields_1,
            "hdr.arp.tpa": match_fields_2
        },
        action_name="MyIngress.send_arp_reply",
        action_params=action_params
    )
    sw.WriteTableEntry(table_entry)

def readTableRules(p4info_helper, sw):
    """
    Reads the table entries from all tables on the switch.
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            # use the p4info_helper to translate the IDs in the entry to names
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
            print()

def printtxt(p4info_helper, insw, esw, counter_name, index):
    for response in insw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            if os.path.exists(f"{insw.name}{esw.name}.txt"):
                with open(f"{insw.name}{esw.name}.txt", "a") as file:
                    file.write("%s->%s\n" % (insw.name, esw.name))
                    file.write("%s sent %d packets, counter ID: %d\n" % (insw.name, counter.data.packet_count, index))
                    file.write("%s received %d packets, counter ID: %d\n" % (esw.name, counter.data.packet_count, index))
            else:
                with open(f"{esw.name}{insw.name}.txt", "a") as file:
                    file.write("%s->%s\n" % (insw.name, esw.name))
                    file.write("%s sent %d packets, counter ID: %d\n" % (insw.name, counter.data.packet_count, index))
                    file.write("%s received %d packets, counter ID: %d\n" % (esw.name, counter.data.packet_count, index))

def printCounter(p4info_helper, sw, counter_name, index):
    """
    Reads the specified counter at the specified index from the switch.
    """
    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            print("%s %s %d: %d packets (%d bytes)" % (
                sw.name, counter_name, index,
                counter.data.packet_count, counter.data.byte_count
            ))

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create switch connections
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Establish master arbitration
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install P4 program
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)

        # Write tunnel rules to s1, s2, s3
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s2, tunnel_id=100,
                        dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2")
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s3, tunnel_id=200,
                        dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3")
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s1, tunnel_id=101,
                        dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1")
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s3, tunnel_id=300,
                        dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3")
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s1, tunnel_id=201,
                        dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1")
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s2, tunnel_id=301,
                        dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2")

        # Write arp_match table entries to s1, s2, s3
        writeArp_matchRules(p4info_helper, s1, [1], ["10.0.1.10", 32], {"macAddr": "08:00:00:00:01:00"})
        writeArp_matchRules(p4info_helper, s2, [1], ["10.0.2.20", 32], {"macAddr": "08:00:00:00:02:00"})
        writeArp_matchRules(p4info_helper, s3, [1], ["10.0.3.30", 32], {"macAddr": "08:00:00:00:03:00"})

        # Read table entries from s1, s2 and s3
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)
        readTableRules(p4info_helper, s3)

        # Print the tunnel counters every 1 second
        while True:
            sleep(1)
            print('\n----- Reading tunnel counters -----')

            print('\n----- s1 -> s2 -----')
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)
            printtxt(p4info_helper, s1, s2, "MyIngress.ingressTunnelCounter", 100)

            print('\n----- s2 -> s1 -----')
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 101)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 101)
            printtxt(p4info_helper, s2, s1, "MyIngress.ingressTunnelCounter", 101)

            print('\n----- s1 -> s3 -----')
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 200)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 200)
            printtxt(p4info_helper, s1, s3, "MyIngress.ingressTunnelCounter", 200)

            print('\n----- s3 -> s1 -----')
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 201)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 201)
            printtxt(p4info_helper, s3, s1, "MyIngress.ingressTunnelCounter", 201)

            print('\n----- s2 -> s3 -----')
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 300)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 300)
            printtxt(p4info_helper, s2, s3, "MyIngress.ingressTunnelCounter", 300)

            print('\n----- s3 -> s2 -----')
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301)
            printtxt(p4info_helper, s3, s2, "MyIngress.ingressTunnelCounter", 301)

            print('\n------------ Finished ------------')

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)

